"""
Author: '夜微凉'
Date: 2025-01-09 10:41:42
LastEditors: '夜微凉'
LastEditTime: 2025-01-09 11:01:26
FilePath: /lsq-net-master/lrg/main.py
Description: 

"""

from utils.util import *
from utils import data_loader
from loguru import logger
from utils import quantize
logger.add("./lsq.log")
def main():
    # 获取配置文件
    args = get_config(r"./config.yaml")
    logger.debug(args.keys())

    # Initialize data loader
    train_loader, test_loader = data_loader.load_data(args["dataloader"])
    logger.info(
        "Dataset `%s` size:" % args["dataloader"]["dataset"]
        + "\n          Training Set = %d (%d)"
        % (len(train_loader.sampler), len(train_loader))
        + "\n              Test Set = %d (%d)"
        % (len(test_loader.sampler), len(test_loader))
    )
    model=get_model_cifar10(args)
    #train(args,model,train_loader,test_loader)
    # 转换量化节点
    #modules_to_replace=quantize.find_modules_to_quantize(model,args["quantize"])
    # logger.info(modules_to_replace)
    #model = quantize.replace_module_by_names(model, modules_to_replace)
    #model=model.cuda()
    train(args,model,train_loader,test_loader)
    # train(train_loader,model)
    # test(test_loader,model)



if __name__ == "__main__":
    main()
